import argparse
import torch
import onnx
import os
from yolov5.models.yolo import Model
from yolov5.utils.torch_utils import select_device
from utils.deployment.trt_export import export_tensorrt


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='runs/train/exp/weights/best.pt', help='model.pt path(s)')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image (h,w)')
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--include', nargs='+', default=['onnx'], help='include formats')
    parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
    parser.add_argument('--simplify', action='store_true', help='simplify onnx model')
    parser.add_argument('--trt', action='store_true', help='export tensorrt engine')
    parser.add_argument('--dynamic', action='store_true', help='dynamic axes')
    return parser.parse_args()


def export_onnx(model, im, file, opset, dynamic, simplify):
    """Export model to ONNX format"""
    try:
        import onnx
        print(f'\nStarting ONNX export with onnx {onnx.__version__}...')
        torch.onnx.export(
            model,
            im,
            file,
            opset_version=opset,
            do_constant_folding=True,
            input_names=['images'],
            output_names=['output'],
            dynamic_axes={'images': {0: 'batch_size'}, 'output': {0: 'batch_size'}} if dynamic else None
        )

        # Check exported model
        onnx_model = onnx.load(file)
        onnx.checker.check_model(onnx_model)

        # Simplify onnx model
        if simplify:
            try:
                import onnxsim
                print(f'Simplifying with onnxsim {onnxsim.__version__}...')
                onnx_model, check = onnxsim.simplify(onnx_model)
                assert check, 'Simplified ONNX model could not be validated'
                onnx.save(onnx_model, file)
            except Exception as e:
                print(f'ONNX simplification failed: {e}')

        print(f'ONNX export success, saved as {file}')
    except Exception as e:
        print(f'ONNX export failure: {e}')


def main(opt):
    opt.img_size *= 2 if len(opt.img_size) == 1 else 1
    device = select_device(opt.device)

    # Load PyTorch model
    model = Model('models/yolov5s.yaml').to(device)
    weights = torch.load(opt.weights, map_location=device)['model'].float()
    model.load_state_dict(weights.state_dict())
    model.eval()

    # Create example input
    im = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device)

    # Export ONNX
    if 'onnx' in opt.include or opt.trt:
        onnx_file = opt.weights.replace('.pt', '.onnx')
        export_onnx(model, im, onnx_file, 12, opt.dynamic, opt.simplify)

    # Export TensorRT
    if opt.trt:
        export_tensorrt(onnx_file, opt.img_size, opt.half)

if __name__ == '__main__':
    opt = parse_opt()
    main(opt)